import datetime
import uuid

from metadata.generated.schema.entity.data.table import (
    Column,
    ColumnName,
    Constraint,
    DataType,
    Table,
    TableData,
    TableType,
)
from metadata.generated.schema.type.basic import (
    EntityName,
    FullyQualifiedEntityName,
    Markdown,
    Uuid,
)
from metadata.generated.schema.type.tagLabel import (
    LabelType,
    State,
    TagFQN,
    TagLabel,
    TagSource,
)
from metadata.ingestion.models.table_metadata import ColumnTag
from metadata.sampler.models import SampleData

table = Table(
    id=Uuid(root=uuid.uuid4()),
    name=EntityName(root="customers"),
    displayName=None,
    fullyQualifiedName=FullyQualifiedEntityName(
        root="Service.database.schema.customers"
    ),
    description=Markdown(root="testing comments here"),
    tableType=TableType.Partitioned,
    columns=[
        Column(
            name=ColumnName(root="customer_id"),
            displayName=None,
            dataType=DataType.INT,
            arrayDataType=None,
            dataLength=1,
            precision=1,
            scale=None,
            dataTypeDisplay="integer",
            description=Markdown(root="the ID of the customer"),
            fullyQualifiedName=FullyQualifiedEntityName(
                root="Service.database.schema.customers.customer_id"
            ),
            constraint=Constraint.NULL,
        ),
        Column(
            name=ColumnName(root="first_name"),
            displayName=None,
            dataType=DataType.VARCHAR,
            arrayDataType=None,
            dataLength=20,
            precision=20,
            scale=None,
            dataTypeDisplay="character varying(20)",
            description=Markdown(root="Customer's first name. PII."),
            fullyQualifiedName=FullyQualifiedEntityName(
                root="Service.database.schema.customers.first_name"
            ),
            constraint=Constraint.NULL,
        ),
        Column(
            name=ColumnName(root="last_name"),
            displayName=None,
            dataType=DataType.VARCHAR,
            arrayDataType=None,
            dataLength=2,
            precision=2,
            scale=None,
            dataTypeDisplay="character varying(2)",
            description=Markdown(root="First letter of the Last name of the customer"),
            fullyQualifiedName=FullyQualifiedEntityName(
                root="Service.database.schema.customers.last_name"
            ),
            constraint=Constraint.NULL,
        ),
        Column(
            name=ColumnName(root="first_order"),
            displayName=None,
            dataType=DataType.DATE,
            arrayDataType=None,
            dataLength=1,
            precision=1,
            scale=None,
            dataTypeDisplay="date",
            description=Markdown(root="Date (UTC) of a customer's first order"),
            fullyQualifiedName=FullyQualifiedEntityName(
                root="Service.database.schema.customers.first_order"
            ),
            constraint=Constraint.NULL,
        ),
        Column(
            name=ColumnName(root="most_recent_order"),
            displayName=None,
            dataType=DataType.DATE,
            arrayDataType=None,
            dataLength=1,
            precision=1,
            scale=None,
            dataTypeDisplay="date",
            description=Markdown(root="Date (UTC) of a customer's most recent order"),
            fullyQualifiedName=FullyQualifiedEntityName(
                root="Service.database.schema.customers.most_recent_order"
            ),
            constraint=Constraint.NULL,
        ),
        Column(
            name=ColumnName(root="number_of_orders"),
            displayName=None,
            dataType=DataType.BIGINT,
            arrayDataType=None,
            dataLength=1,
            precision=1,
            scale=None,
            dataTypeDisplay="bigint",
            description=Markdown(
                root="Count of the number of orders a customer has placed"
            ),
            fullyQualifiedName=FullyQualifiedEntityName(
                root="Service.database.schema.customers.number_of_orders"
            ),
            constraint=Constraint.NULL,
        ),
        Column(
            name=ColumnName(root="customer_lifetime_value"),
            displayName=None,
            dataType=DataType.BIGINT,
            arrayDataType=None,
            dataLength=1,
            precision=1,
            scale=None,
            dataTypeDisplay="bigint",
            description=None,
            fullyQualifiedName=FullyQualifiedEntityName(
                root="Service.database.schema.customers.customer_lifetime_value"
            ),
            constraint=Constraint.NULL,
        ),
    ],
)

sample_data = SampleData(
    data=TableData(
        columns=[
            ColumnName(root="customer_id"),
            ColumnName(root="first_name"),
            ColumnName(root="last_name"),
            ColumnName(root="first_order"),
            ColumnName(root="most_recent_order"),
            ColumnName(root="number_of_orders"),
            ColumnName(root="customer_lifetime_value"),
        ],
        rows=[
            [5, "Katherine", "R.", None, None, None, None],
            [
                7,
                "Martin",
                "M.",
                datetime.date(2018, 1, 14),
                datetime.date(2018, 1, 14),
                1,
                26,
            ],
            [
                9,
                "Jennifer",
                "F.",
                datetime.date(2018, 3, 17),
                datetime.date(2018, 3, 17),
                1,
                30,
            ],
            [
                11,
                "Fred",
                "S.",
                datetime.date(2018, 3, 23),
                datetime.date(2018, 3, 23),
                1,
                3,
            ],
            [
                13,
                "Kathleen",
                "M.",
                datetime.date(2018, 3, 7),
                datetime.date(2018, 3, 7),
                1,
                26,
            ],
            [15, "Teresa", "H.", None, None, None, None],
            [17, "Kimberly", "R.", None, None, None, None],
            [
                19,
                "Virginia",
                "F.",
                datetime.date(2018, 3, 16),
                datetime.date(2018, 3, 16),
                1,
                3,
            ],
            [
                21,
                "Willie",
                "H.",
                datetime.date(2018, 3, 28),
                datetime.date(2018, 3, 28),
                1,
                22,
            ],
            [23, "Mildred", "A.", None, None, None, None],
            [
                25,
                "Victor",
                "H.",
                datetime.date(2018, 1, 17),
                datetime.date(2018, 3, 20),
                2,
                24,
            ],
            [
                27,
                "Benjamin",
                "B.",
                datetime.date(2018, 2, 21),
                datetime.date(2018, 4, 4),
                2,
                27,
            ],
            [29, "Benjamin", "K.", None, None, None, None],
            [
                31,
                "Jane",
                "G.",
                datetime.date(2018, 2, 17),
                datetime.date(2018, 2, 17),
                1,
                18,
            ],
            [
                33,
                "Katherine",
                "M.",
                datetime.date(2018, 2, 13),
                datetime.date(2018, 2, 13),
                1,
                14,
            ],
            [
                35,
                "Sara",
                "T.",
                datetime.date(2018, 2, 21),
                datetime.date(2018, 3, 21),
                2,
                34,
            ],
            [37, "Shirley", "J.", None, None, None, None],
            [
                39,
                "Louise",
                "W.",
                datetime.date(2018, 1, 18),
                datetime.date(2018, 1, 18),
                1,
                10,
            ],
            [
                41,
                "Gloria",
                "C.",
                datetime.date(2018, 4, 7),
                datetime.date(2018, 4, 7),
                1,
                10,
            ],
            [43, "Kelly", "N.", None, None, None, None],
            [45, "Scott", "B.", None, None, None, None],
            [
                47,
                "Marie",
                "P.",
                datetime.date(2018, 3, 26),
                datetime.date(2018, 3, 31),
                2,
                36,
            ],
            [49, "Judy", "N.", None, None, None, None],
            [
                51,
                "Howard",
                "R.",
                datetime.date(2018, 1, 28),
                datetime.date(2018, 2, 23),
                3,
                99,
            ],
            [
                53,
                "Anne",
                "B.",
                datetime.date(2018, 1, 12),
                datetime.date(2018, 3, 11),
                2,
                39,
            ],
            [55, "Nicholas", "R.", None, None, None, None],
            [
                57,
                "Paul",
                "W.",
                datetime.date(2018, 1, 31),
                datetime.date(2018, 3, 5),
                2,
                26,
            ],
            [
                59,
                "Adam",
                "A.",
                datetime.date(2018, 1, 15),
                datetime.date(2018, 1, 15),
                1,
                1,
            ],
            [61, "Timothy", "R.", None, None, None, None],
            [
                63,
                "Edward",
                "G.",
                datetime.date(2018, 3, 3),
                datetime.date(2018, 4, 3),
                2,
                23,
            ],
            [65, "Brenda", "W.", None, None, None, None],
            [67, "Michael", "H.", None, None, None, None],
            [
                69,
                "Janet",
                "P.",
                datetime.date(2018, 2, 2),
                datetime.date(2018, 3, 18),
                2,
                32,
            ],
            [
                71,
                "Gerald",
                "C.",
                datetime.date(2018, 1, 18),
                datetime.date(2018, 2, 24),
                3,
                44,
            ],
            [73, "Alan", "B.", None, None, None, None],
            [75, "Andrea", "H.", None, None, None, None],
            [77, "Anne", "W.", None, None, None, None],
            [
                79,
                "Jack",
                "R.",
                datetime.date(2018, 2, 28),
                datetime.date(2018, 3, 11),
                2,
                27,
            ],
            [81, "Shirley", "H.", None, None, None, None],
            [83, "Virginia", "R.", None, None, None, None],
            [
                85,
                "Theresa",
                "M.",
                datetime.date(2018, 2, 10),
                datetime.date(2018, 4, 9),
                2,
                33,
            ],
            [87, "Phillip", "B.", None, None, None, None],
            [
                89,
                "Margaret",
                "J.",
                datetime.date(2018, 4, 7),
                datetime.date(2018, 4, 7),
                1,
                14,
            ],
            [
                91,
                "Todd",
                "W.",
                datetime.date(2018, 3, 27),
                datetime.date(2018, 3, 27),
                1,
                29,
            ],
            [
                93,
                "Frances",
                "R.",
                datetime.date(2018, 3, 1),
                datetime.date(2018, 3, 1),
                1,
                2,
            ],
            [95, "Lisa", "P.", None, None, None, None],
            [97, "Shirley", "D.", None, None, None, None],
            [
                99,
                "Mary",
                "G.",
                datetime.date(2018, 1, 14),
                datetime.date(2018, 2, 14),
                2,
                44,
            ],
            [
                1,
                "Michael",
                "P.",
                datetime.date(2018, 1, 1),
                datetime.date(2017, 12, 31),
                2,
                33,
            ],
            [
                3,
                "Kathleen",
                "P.",
                datetime.date(2018, 1, 2),
                datetime.date(2018, 1, 1),
                3,
                65,
            ],
        ],
    )
)

expected_column_tags = [
    ColumnTag(
        column_fqn="Service.database.schema.customers.first_name",
        tag_label=TagLabel(
            source=TagSource.Classification,
            labelType=LabelType.Generated,
            state=State.Suggested,
            name="Sensitive",
            tagFQN=TagFQN(
                root="PII.Sensitive",
            ),
        ),
    ),
    ColumnTag(
        column_fqn="Service.database.schema.customers.first_order",
        tag_label=TagLabel(
            source=TagSource.Classification,
            labelType=LabelType.Generated,
            state=State.Suggested,
            name="NonSensitive",
            tagFQN=TagFQN(
                root="PII.NonSensitive",
            ),
        ),
    ),
    ColumnTag(
        column_fqn="Service.database.schema.customers.most_recent_order",
        tag_label=TagLabel(
            source=TagSource.Classification,
            labelType=LabelType.Generated,
            state=State.Suggested,
            name="NonSensitive",
            tagFQN=TagFQN(
                root="PII.NonSensitive",
            ),
        ),
    ),
]
